package com.example.aidemo.controller;

import com.example.aidemo.converter.Markdown2Html;
import jakarta.servlet.http.HttpServletRequest;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.model.Media;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.OpenAiImageModel;
import org.springframework.ai.openai.OpenAiImageOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.net.MalformedURLException;
import java.net.URL;
import java.util.List;
import java.util.Map;

@RestController
public class ChatDemoController {

    @Autowired
    private ChatClient chatClient;
    @Autowired
    private OpenAiChatModel openAiChatModel;
    @Autowired
    private OpenAiImageModel imageModel;
    @Autowired
    private ChatMemory chatMemory;

    /**
     * openAi 聊天
     *
     * @param message
     * @return
     */
    @GetMapping("/ai/openAiChat")
    public Map openAiChat(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        return Map.of("generation", openAiChatModel.call(message));
    }

    /**
     * ChatClient使用(流式调用)
     *
     * @param message
     * @param voice
     * @return
     */
    @GetMapping("/ai/chatClient")
    Map<String, String> completion(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message,
                                   String voice) {
        return Map.of("completion", chatClient.prompt()
                .system(sp -> sp.param("voice", voice))
                .user(message)
                .call()
                .content());
    }

    /**
     * 带有聊天记录功能的 ChatClient使用(流式调用)
     *
     * @param request
     * @param message
     * @return
     */
    @GetMapping("/ai/chatMemory")
    String generationByChatMemory(HttpServletRequest request, @RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        String sessionId = request.getSession().getId();
        chatMemory.add(sessionId, new UserMessage(message));
        String content = this.chatClient.prompt()
                .advisors(new MessageChatMemoryAdvisor(chatMemory))
                .user(message)
                .call()
                .content();
        String htmlContent = Markdown2Html.markdown2Html(content);
        chatMemory.add(sessionId, new AssistantMessage(htmlContent));
        return htmlContent;
    }

    /**
     * 图片生成(文生图)
     *
     * @param message
     * @return
     */
    @GetMapping("/ai/imageGeneration")
    public Map imageGeneration(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        OpenAiImageOptions imageOptions = OpenAiImageOptions.builder()
                .withQuality("hd")
                .withN(1)
                .withHeight(1024)
                .withWidth(1024).build();
        ImagePrompt imagePrompt = new ImagePrompt(message, imageOptions);
        ImageResponse response = imageModel.call(imagePrompt);
        return Map.of("generation", response.getResult().getOutput().getUrl());
    }

    /**
     * 多模态
     *
     * @param message
     * @return
     * @throws MalformedURLException
     */
    @GetMapping("/ai/multimodal")
    public String Multimodal(@RequestParam(value = "message", defaultValue = "解释一下你在这张图片上看到了什么?") String message) throws MalformedURLException {
        var userMessage = new UserMessage(message,
                List.of(new Media(MimeTypeUtils.IMAGE_PNG,
                        new URL("https://docs.spring.io/spring-ai/reference/1.0-SNAPSHOT/_images/multimodal.test.png"))));

        ChatResponse response = openAiChatModel.call(new Prompt(List.of(userMessage),
                OpenAiChatOptions.builder().withModel(OpenAiApi.ChatModel.GPT_4_O.getValue()).build()));
        return response.getResult().getOutput().getContent();
    }
}
